package me.aias.example;

import ai.djl.ModelException;
import ai.djl.inference.Predictor;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.training.util.ProgressBar;
import ai.djl.translate.TranslateException;
import me.aias.example.utils.LacTranslator;
import me.aias.example.utils.TranslationTranslator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.nio.file.Path;
import java.util.Arrays;
import java.util.Locale;

/**
 * 中文翻译为英文
 *
 * @author calvin
 * @mail 179209347@qq.com
 * @website www.aias.top
 */
public final class TranslationExample {

    private static final Logger logger = LoggerFactory.getLogger(TranslationExample.class);

    public static void main(String[] args) throws IOException, TranslateException, ModelException {
        // 分词
        Criteria<String, String[][]> lacCriteria = Criteria.builder()
                .setTypes(String.class, String[][].class)
                .optModelUrls("https://aias-home.oss-cn-beijing.aliyuncs.com/models/nlp_models/lac.zip")
//                .optModelPath(Path.of("E:/machine_translation/translation_zh_en_sdk/src/main/resources/lac.zip"))
                .optTranslator(new LacTranslator())
                .optProgress(new ProgressBar())
                .optEngine("PaddlePaddle")
                .build();
        // 翻译
        Criteria<String[], String[]> sentaCriteria = Criteria.builder()
                .setTypes(String[].class, String[].class)
                .optModelUrls("https://aias-home.oss-cn-beijing.aliyuncs.com/models/nlp_models/translation/translation_zh_en.zip")
//                .optModelPath(Path.of("E:/machine_translation/translation_zh_en_sdk/src/main/resources/translation_zh_en.zip"))
                .optTranslator(new TranslationTranslator())
                .optEngine("PaddlePaddle")
                .optModelName("inference")
                .optProgress(new ProgressBar())
                .build();

        try (ZooModel<String, String[][]> lacModel = lacCriteria.loadModel();
             Predictor<String, String[][]> lacPredictor = lacModel.newPredictor();
             ZooModel<String[], String[]> sentaModel = sentaCriteria.loadModel();
             Predictor<String[], String[]> sentaPredictor = sentaModel.newPredictor()) {

            String input = "我们经常在工作中会将一些文本进行多语言翻译";
            logger.info("输入句子: {}", input);

            String[][] lacResult = lacPredictor.predict(input);
            // 分词
            logger.info("Words : " + Arrays.toString(lacResult[0]));
            // 词性
            logger.info("Tags : " + Arrays.toString(lacResult[1]));

            // 翻译结果
            String[] translationResult = sentaPredictor.predict(lacResult[0]);
            for (int i = 0; i < translationResult.length; i++) {
                logger.info("T" + i + ": " + translationResult[i]);
            }
            logger.info(Arrays.toString(translationResult));
        }
    }
}
